import math

class TreeNode(): # Tree data structure
    def __init__(self, label, value) -> None:
        self.label = label
        self.value = value
        self.layer = None
        self.huffman_code = None
        self.children = []
    def add_child(self, child):
        self.children.append(child)

# source = [['x_1', 0.2],
#           ['x_3', 0.18],
#           ['x_4', 0.17],
#           ['x_2', 0.19],
#           ['x_5', 0.15],
#           ['x_6', 0.1],
#           ['x_7', 0.01]]

# source = [['1', 0.25], 
#           ['2', 0.125], 
#           ['-2', 0.0625], 
#           ['-1', 0.1875], 
#           ['0', 0.25], 
#           ['-3', 0.0625], 
#           ['3', 0.0625]]

class Huffuman_encode_2():
    # Binary Huffman encoding
    def __init__(self, source, SUM) -> None:
        self.source = source
        self.check_right_probability()
        self.check_label_length()
        self.dict_source = dict(self.source)
        self.build_nodes()
        self.root = None # Huffman tree root
        self.Huffuman_code = None
        self.SUM_ALL = SUM

        self.Huffuman_tree()
        self.encoding()
    def check_label_length(self):
        Source = self.source
        MAX_len = max([len(item[0]) for item in Source])
        self.MAX_label_len = MAX_len
        for ITEM in Source:
            if(len(ITEM[0]) < MAX_len):
                label = ('*' * (MAX_len - len(ITEM[0]))) + ITEM[0]
                ITEM[0] = label
        self.source = Source
    def check_right_probability(self):
        self.probability = [x[1] for x in self.source]
        sum = 0
        for value in self.probability:
            sum += value
        if(sum == 1):
            pass
        else:
            print("The probability is not right (Sum ≠ 1). Please try again.")
            exit()

    def sort_tree_nodes(self, tree_nodes):
        # Sort the nodes list based on probability.
        return sorted(tree_nodes,
                      key=lambda x: x.value,
                      reverse=True)

    def build_nodes(self):
        # Build nodes list
        nodes_list = []
        for value in self.source:
            node = TreeNode(value[0], value[1])
            nodes_list.append(node)
        self.nodes_list = nodes_list
    def Huffuman_tree(self):
        # Build Huffman tree.
        nodes_list = self.nodes_list
        size = len(nodes_list)
        while(size > 1):
            nodes_list = self.sort_tree_nodes(nodes_list)
            child1 = nodes_list.pop()
            child2 = nodes_list.pop()
            parent_label = child1.label + ' + ' + child2.label
            parent_value = child1.value + child2.value
            parent = TreeNode(parent_label, parent_value)
            parent.add_child(child2)
            parent.add_child(child1)
            nodes_list.append(parent)
            size -= 1
        self.root = nodes_list[0]
    
    def encoding(self):
        # The following two are sub functions.
        def Ensure_Huffman_code(node, code, layer):
            # Traverse the tree. (Deep first traversal)
            node.huffman_code = code
            node.layer = layer
            # print(node.label, ' ', node.value, ' ', node.layer, ' ', node.huffman_code)
            i = 0
            for child in node.children:
                Ensure_Huffman_code(child, code + str(i), layer+1)
                i += 1

        def Choose_source_node(node):
            if(len(node.label) == self.MAX_label_len): # If it is a source node(like 'x_1').
                temp_list = []
                temp_list.append(node.label)
                temp_list.append(dict_source[node.label])
                temp_list.append(node.layer)
                temp_list.append(node.huffman_code)
                Huffuman_code.append(temp_list)
            for child in node.children:
                Choose_source_node(child)

        # Encoding
        Ensure_Huffman_code(self.root, '', 0)
        dict_source = self.dict_source
        Huffuman_code = [] # [label, num, length, Huffman code]
        Choose_source_node(self.root)
        Huffuman_code = sorted(Huffuman_code, key=lambda x: x[2], reverse=False)
        index = 1
        for Item in Huffuman_code:
            Item.insert(1, int(Item[1] * self.SUM_ALL))
            Item.pop(2)
            Item[0] = Item[0].replace('*', '')
            Item.insert(0, index)
            index += 1
        self.Huffuman_code = Huffuman_code

    def print_results(self):
        print("label, num, code length, Huffman code")
        for item in self.Huffuman_code:
            print(item)

# huffman2 = Huffuman_encode_2(source, 16)
# huffman2.print_results()

'''
huffman2 = Huffuman_encode_2(source)
huffman2.print_results()

Output:
label, probability, code length, Huffman code
['x_1', 0.2, 2, '10']
['x_2', 0.19, 2, '11']
['x_3', 0.18, 3, '000']
['x_4', 0.17, 3, '001']
['x_5', 0.15, 3, '010']
['x_6', 0.1, 4, '0110']
['x_7', 0.01, 4, '0111']

Entropy = 2.608682761140783
Average encoding symbol length(K) = 2.72
Average transmission efficiency(R) = 0.9590745445370524
'''




